/** @file admm_arguments.h
 *
 *  @author Dongryeol Lee (drselee@gmail.com)
 */

#ifndef MLPACK_ADMM_ADMM_ARGUMENTS_H
#define MLPACK_ADMM_ADMM_ARGUMENTS_H

#include "core/abstract/abstract_arguments.h"
#include "core/data/matrix_dataset.h"
#include "core/util/random.h"
#include "mlpack/admm/admm_result.h"

namespace mlpack {
namespace admm {

class ADMMArguments : public core::abstract::AbstractArguments {
  public:

    /** @brief The rho parameter in ADMM.
     */
    double admm_rho_;

    /** @brief The inverse of the rho parameter in ADMM.
     */
    double admm_rho_inverse_;

    /** @brief The squared norm of each reference point.
     */
    arma::vec squared_norm_reference_point_;

    /** @brief The SVM C parameter.
     */
    double svm_c_parameter_;

  private:

    bool Validate_() const {
      if(vm_["admm_rho_in"].as<double>() <= 0.0) {
        std::cerr << "The --admm_rho_in requires a positive real number.\n";
        return false;
      }
      if(vm_["svm_c_parameter_in"].as<double>() <= 0.0) {
        std::cerr <<
                  "The --svm_c_parameter_in requires a positive real number.\n";
        return false;
      }
      return true;
    }

  public:

    double squared_norm_reference_point(int point_id) const {
      return squared_norm_reference_point_[ point_id ];
    }

    /** @brief The default constructor.
     */
    ADMMArguments() {
      admm_rho_ = 1.0;
      admm_rho_inverse_ = 1.0 / admm_rho_;
      svm_c_parameter_ = 1.0;

      // Add more descriptions.
      desc_->add_options()(
        "admm_rho_in",
        boost::program_options::value<double>()->default_value(0.1),
        "OPTIONAL The rho value that controls the convergence in the "
        "ADMM framework."
      )(
        "loss_in",
        boost::program_options::value<std::string>()->default_value(
          "squared_hinge_loss"),
        "OPTIONAL The type of loss function to use. One of: "
        "  hinge_loss, squared_hinge_loss"
      )(
        "regularization_in",
        boost::program_options::value<std::string>()->default_value("l2"),
        "OPTIONAL The type of regularization to use. One of: "
        "  l1, l2"
      )(
        "svm_c_parameter_in",
        boost::program_options::value<double>()->default_value(20.0),
        "OPTIONAL The SVM C parameter value for L1 and L2 SVM."
      );
    }

    /** @brief The destructor.
     */
    ~ADMMArguments() {
    }

    bool Parse_(
      boost::mpi::communicator & world,
      const std::vector< std::string > & args) {

      // Validate the arguments. Only immediate termination is allowed
      // here, the parsing is done later.
      if(! this->Validate_()) {
        return false;
      }

      // end of argument validation.

      // Parse the rho value for the ADMM.
      admm_rho_ = vm_["admm_rho_in"].as<double>();
      admm_rho_inverse_ = 1.0 / admm_rho_;

      // From the reference set, pre-compute the squared norm of each
      // reference point.
      squared_norm_reference_point_.zeros(
        reference_table_->num_points());
      for(int i = 0; i < reference_table_->num_points(); i++) {
        const core::data::alias_vec *targets = NULL;
        const core::data::alias_vec &point =
          reference_table_->get_instance(i, &targets);
        squared_norm_reference_point_[i] =
          arma::dot(point, point);
      }

      // Parse the maximum number of iterations.
      max_num_iterations_ = vm_["max_num_iterations_in"].as<int>();
      if(world.rank() == 0) {
        std::cout << "Will run for " <<
                  max_num_iterations_ << " rounds.\n";
      }

      // Parse the model file name.
      model_out_ = vm_["model_out"].as<std::string>();

      // Parse the prediction output file name.
      predictions_out_ = vm_["predictions_out"].as<std::string>();

      // Parse the number of threads.
      num_threads_ = vm_["num_threads_in"].as<int>();
      if(world.rank() == 0) {
        std::cout << "Using " << num_threads_ <<
                  " threads for shared memory parallelism...\n";
      }

      // Parse the SVM C parameter.
      svm_c_parameter_ = vm_["svm_c_parameter_in"].as<double>();
      return false;
    }
};
}
}

#endif
